# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Ring-Polymer PIMD integrator
"""

from typing import Tuple

import torch
import math
from torch import Tensor

from .integrator import Integrator, _integrator_register
from ..thermostat import Thermostat
from ..barostat import Barostat
from ..constraint import Constraint
from ...system import Molecule
from ...function import get_arguments, get_tensor


@_integrator_register('rpmd')
class RPMD(Integrator):
    r"""A leap-frog integrator based on "middle scheme" developed by Jian Liu, et al. It is a subclass of `Integrator`.

    Reference:

        Zhang, Z.; Yan, K; Liu, X.; Liu, J.
        A Leap-Frog Algorithm-based Efficient Unified Thermostat Scheme for Molecular Dynamics [J].
        Chinese Science Bulletin, 2018, 63(33): 3467-3483.

    Args:

        system (Molecule):          Simulation system
        constraint (Constraint):    Constraint algorithm. Default: None

    """
    def __init__(self,
                 system: Molecule,
                 temperature = 300.,
                 gamma0 = 10.,
                 remove_com = True,
                 method = 'BAOAB', 
                 **kwargs
                 ):

        super().__init__(
            system=system,
            thermostat=None,
            barostat=None,
            constraint=None,
        )
        self._kwargs = get_arguments(locals(), kwargs)
        self.name = 'PILE'
        n = self.num_walker
        C_jk = get_tensor(torch.zeros((n,n)), torch.float32)
        for j in range(n):
            for k in range(n):
                if k ==0:
                    C_jk[j][k] = math.sqrt(1/n) 
                elif 1<=k<=(n/2-1):
                    C_jk[j][k] = math.sqrt(2/n)*math.cos(2*math.pi*(j+1)*k/n)
                elif k == n/2:
                    C_jk[j][k] = math.sqrt(1/n)*(-1)**(j+1)
                elif n/2+1<=k<=n-1:
                    C_jk[j][k] = math.sqrt(2/n)*math.sin(2*math.pi*(j+1)*k/n)
        self.C_jk = C_jk

        self.ref_temp = self._get_mw_tensor(temperature, 'temperature')
        boltzmann = system.units.boltzmann
        hbar = 1.054572e-37 * 6.02e23 * 1e12
        self.beta = 1. / (boltzmann * self.ref_temp)
        self.beta_n = self.beta / n
        omega_n = 1.0 / (self.beta_n * hbar)
        omega_k = 2 * omega_n * torch.sin(get_tensor(torch.arange(n).reshape(n,1,1)*torch.pi/n, torch.float32)) + 1e-8
        
        self.a = torch.cos(omega_k*0.5*self.time_step)
        self.b = torch.sin(omega_k*0.5*self.time_step)/omega_k
        self.c = -torch.sin(omega_k*0.5*self.time_step)*omega_k
        self.d = torch.cos(omega_k*0.5*self.time_step)

        self.gamma = 2*omega_k
        self.gamma0 = gamma0
        self.gamma[0] = self.gamma0
        self.c1 = torch.exp(-self.gamma*self.time_step)
        self.c2 = torch.sqrt(1-self.c1**2) / self.kinetic_unit_scale
        
        self.sqrt_inv_mass= torch.sqrt(self._inv_mass)
        self.remove_com = remove_com
    
    def print_info(self):
        info = {'Simulation Ensemble': 'NVT Ensemble',
                'Integrator Type': 'RPMD',
                'time step': self.time_step,
                'temperature': self.temperature.item(),
                'gamma0': self.gamma0,}
        return info

    def forward(self,
                coordinate: Tensor,
                velocity: Tensor,
                force: Tensor,
                energy: Tensor,
                virial: Tensor = None,
                pbc_box: Tensor = None,
                steps: int = 0,
                **kwargs
                  ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:

        # (B,A,D) = (B,A,D) * (B,A,1)
        acceleration = self.acc_unit_scale * force * self._inv_mass
        # v(t+0.5) = v(t-0.5) + a(t) * dt
        velocity_half = velocity + acceleration * self.time_step
        # return velocity_half, velocity, acceleration, self.time_step
        # (B,A,D) = (B,A,D) - (B,1,D)
        if self.remove_com:
            velocity_half = velocity_half - self.get_com_velocity(velocity_half)

        v_tilde = torch.einsum("jk,jad->kad",self.C_jk,velocity_half)
        x_tilde = torch.einsum("jk,jad->kad",self.C_jk,coordinate)
        # R(t+0.5) = R(t) + v(t+0.5) * dt
        
        x_tilde_half = self.a * x_tilde + self.b * v_tilde
        v_tilde_half = self.c * x_tilde + self.d * v_tilde

        if self.ref_temp is not None:
            # v'(t+0.5) = f_T[v(t+0.5)]
            v_tilde_half = self.langevin_thermostat(v_tilde_half, steps)
        
        # return v_tilde_half

        # R(t+1) = R(t+0.5) + v'(t+0.5) * dt
        x_tilde_new = self.a * x_tilde_half + self.b * v_tilde_half
        v_tilde_new = self.c * x_tilde_half + self.d * v_tilde_half

        coordinate_new = torch.einsum("jk,kad->jad",self.C_jk,x_tilde_new)
        velocity_half = torch.einsum("jk,kad->jad",self.C_jk,v_tilde_new)

        variables = {'coordinate': coordinate_new,
            'velocity': velocity_half,
            'force': force,
            'energy': energy,
            'virial': virial,
            'pbc_box': pbc_box,
            }
        
        return variables
    
    def langevin_thermostat(self, velocity, step):
        if self.control_step == 1 or step % self.control_step == 0:
            velocity = self.c1 * velocity + self.c2 * torch.sqrt(1 / self.beta_n) * self.sqrt_inv_mass * torch.randn(velocity.shape, device=velocity.device)
        return velocity

    def get_com_momentum(self, momentum:Tensor) -> Tensor:
        # (B,A,D) = (B,A,D) - (B,1,D)
        com_momentum = momentum.sum(-2,keepdim=True)/self._atom_mass.sum(-2,keepdim=True) * self._atom_mass
        return com_momentum

